{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis notebook for Matrix Multiplication benchmark\n",
    "\n",
    "This notebook generates plots and tables from the results of the matrix multiplication benchmark.  See `benchmarks/matrix_multiplication` for more information on the benchmark itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dataclasses\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "import pickle\n",
    "from plotly import express as ex\n",
    "import scipy.io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from symforce.benchmarks.matrix_multiplication.generate_matrix_multiplication_benchmark import (\n",
    "    get_matrices,\n",
    ")\n",
    "\n",
    "matrices = [m[0] for m in get_matrices()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This file should be generated by `run_benchmarks.py`\n",
    "from symforce.benchmarks.run_benchmarks import MatmulBenchmarkConfig\n",
    "\n",
    "with (Path(\"../../..\") / \"benchmark_outputs\" / \"matrix_multiplication_benchmark_results.pkl\").open(\n",
    "    \"rb\"\n",
    ") as f:\n",
    "    d = pickle.load(f)\n",
    "\n",
    "# fmt: off\n",
    "# d = {('b1_ss', 'double', 'sparse', 7, 7, 49, 15): [1110.422429, 14541544066.0, 4528751448.0], ('b1_ss', 'double', 'flattened', 7, 7, 49, 15): [9.130341, 90817350.0, 18226894.0], ('b1_ss', 'double', 'dense_dynamic', 7, 7, 49, 15): [371.027553, 5127401265.0, 1247404369.0], ('b1_ss', 'double', 'dense_fixed', 7, 7, 49, 15): [68.14021, 678885373.0, 419247643.0], ('Tina_DisCog', 'double', 'sparse', 11, 11, 121, 48): [2845.366436, 37224986078.0, 12096487578.0], ('Tina_DisCog', 'double', 'flattened', 11, 11, 121, 48): [33.045554, 470848994.0, 108236516.0], ('Tina_DisCog', 'double', 'dense_dynamic', 11, 11, 121, 48): [496.618231, 6695699808.0, 1805495380.0], ('Tina_DisCog', 'double', 'dense_fixed', 11, 11, 121, 48): [343.571566, 4985289199.0, 1609370025.0], ('n3c4_b2', 'double', 'sparse', 20, 15, 300, 60): [3161.323538, 40365426533.0, 13588919750.0], ('n3c4_b2', 'double', 'flattened', 20, 15, 300, 60): [45.728614, 440863375.0, 92240806.0], ('n3c4_b2', 'double', 'dense_dynamic', 20, 15, 300, 60): [763.965276, 9962125492.0, 2711624958.0], ('n3c4_b2', 'double', 'dense_fixed', 20, 15, 300, 60): [524.794466, 7432520426.0, 2252440577.0], ('bibd_9_3', 'double', 'sparse', 36, 84, 3024, 252): [214.708183, 2409564974.0, 961313801.0], ('bibd_9_3', 'double', 'flattened', 36, 84, 3024, 252): [65.843983, 733343012.0, 311912384.0], ('bibd_9_3', 'double', 'dense_dynamic', 36, 84, 3024, 252): [131.876468, 1579313610.0, 514380660.0], ('bibd_9_3', 'double', 'dense_fixed', 36, 84, 3024, 252): [128.690005, 1482751252.0, 505146003.0], ('lp_sc105', 'double', 'sparse', 105, 163, 17115, 340): [159.657506, 1918261062.0, 739866228.0], ('lp_sc105', 'double', 'flattened', 105, 163, 17115, 340): [40.520398, 460701824.0, 196825221.0], ('lp_sc105', 'double', 'dense_dynamic', 105, 163, 17115, 340): [2074.208572, 18162980911.0, 6383131613.0], ('lp_sc105', 'double', 'dense_fixed', 105, 163, 17115, 340): None}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_dataframe(d, worst_cpu=None):\n",
    "    # Filter out tests with no results\n",
    "    d = {k: v for k, v in d.items() if v is not None}\n",
    "\n",
    "    # Scale relative to worst result\n",
    "    if worst_cpu is not None:\n",
    "        d = {k: [x / worst_cpu[k.matrix_name] for x in v] for k, v in d.items()}\n",
    "\n",
    "    # Remap names\n",
    "    remap = {\n",
    "        \"sparse\": \"Sparse\",\n",
    "        \"dense_dynamic\": \"Dense Dynamic\",\n",
    "        \"dense_fixed\": \"Dense Fixed\",\n",
    "        \"flattened\": \"SymForce\",\n",
    "    }\n",
    "\n",
    "    d = {dataclasses.replace(k, test=remap[k.test]): v for k, v in d.items()}\n",
    "\n",
    "    df = pd.DataFrame(\n",
    "        [dataclasses.astuple(key) + tuple(value) for key, value in d.items()],\n",
    "        columns=(\"mat\", \"scalar\", \"test\", \"M\", \"N\", \"MN\", \"nnz\", \"cpu\", \"instr\", \"l1-load\"),\n",
    "    )\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "absolute_df = make_dataframe(d)\n",
    "absolute_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def summary_stats(df):\n",
    "    tests = df[df.scalar == \"double\"].test.unique()\n",
    "\n",
    "    best_cpu = {}\n",
    "    worst_cpu = {}\n",
    "    shapes = []\n",
    "    for matrix in matrices:\n",
    "        cs = []\n",
    "        for test in tests:\n",
    "            try:\n",
    "                cpu = df.cpu[\n",
    "                    (df.scalar == \"double\") & (df.mat == matrix) & (df.test == test)\n",
    "                ].values[0]\n",
    "            except IndexError:\n",
    "                cpu = \"N/A\"\n",
    "            cs.append(cpu)\n",
    "        best_cpu[matrix] = min([c for c in cs if not isinstance(c, str)])\n",
    "        worst_cpu[matrix] = max([c for c in cs if not isinstance(c, str)])\n",
    "\n",
    "        M = df.M[(df.mat == matrix)].values[0]\n",
    "        N = df.N[(df.mat == matrix)].values[0]\n",
    "        nnz = df.nnz[(df.mat == matrix)].values[0]\n",
    "        shape = f\"{M}x{N}, {round(nnz / (M * N) * 100)}\\\\%\"\n",
    "        shapes.append(shape)\n",
    "\n",
    "    return tests, best_cpu, worst_cpu, shapes\n",
    "\n",
    "\n",
    "tests, best_cpu, worst_cpu, shapes = summary_stats(absolute_df)\n",
    "df = make_dataframe(d, worst_cpu)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ex.bar(\n",
    "    df[df.scalar == \"double\"],\n",
    "    x=\"mat\",\n",
    "    y=\"cpu\",\n",
    "    color=\"test\",\n",
    "    barmode=\"group\",\n",
    "    labels={\"test\": \"Method\", \"mat\": \"Matrix\", \"cpu\": \"Relative CPU %\"},\n",
    "    category_orders={\"test\": [\"Sparse\", \"Dense Dynamic\", \"Dense Fixed\", \"SymForce\"]},\n",
    "    color_discrete_map={\n",
    "        \"Sparse\": ex.colors.qualitative.Plotly[0],\n",
    "        \"Dense Dynamic\": ex.colors.qualitative.Plotly[3],\n",
    "        \"Dense Fixed\": ex.colors.qualitative.Plotly[2],\n",
    "        \"SymForce\": ex.colors.qualitative.Plotly[1],\n",
    "    },\n",
    ")\n",
    "fig.update_layout(\n",
    "    yaxis=dict(tickmode=\"array\", tickvals=[0, 0.5, 1.0], ticktext=[\"0%\", \"50%\", \"100%\"]),\n",
    "    xaxis=dict(\n",
    "        tickmode=\"array\",\n",
    "        tickvals=[0, 1, 2, 3, 4, 5],\n",
    "        ticktext=[\"(a)\", \"(b)\", \"(c)\", \"(d)\", \"(e)\", \"(f)\"],\n",
    "    ),\n",
    ")\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the latex-formatted results table\n",
    "print(\"& \" + \" & \".join(matrices).replace(\"_\", \"\\\\_\") + \" \\\\\\\\\")\n",
    "print(\"& \" + \" & \".join(shapes) + \" \\\\\\\\ \\\\hline\")\n",
    "for test in tests:\n",
    "    s = [test.replace(\"_\", \" \").title()]\n",
    "    for matrix in matrices:\n",
    "        try:\n",
    "            cpu = absolute_df.cpu[\n",
    "                (absolute_df.scalar == \"double\")\n",
    "                & (absolute_df.mat == matrix)\n",
    "                & (absolute_df.test == test)\n",
    "            ].values[0]\n",
    "        except IndexError:\n",
    "            cpu = \"N/A\"\n",
    "        s.append(cpu)\n",
    "\n",
    "    def format(c):\n",
    "        if c in best_cpu.values():\n",
    "            c = round(c * 1e6, 1)  # ms to ns\n",
    "            return f\"\\\\textbf{{{c}}}\"\n",
    "        else:\n",
    "            if isinstance(c, str):\n",
    "                return c\n",
    "            c = round(c * 1e6, 1)  # ms to ns\n",
    "            return f\"{c}\"\n",
    "\n",
    "    s = [format(c) for c in s]\n",
    "    print(\" & \".join(s) + \" \\\\\\\\\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "c70be84c1c0ff77238398e7c58f6c2d1de1437da897bd23db80b2e35f10eba84"
  },
  "kernelspec": {
   "display_name": "symforce-paper",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
